package com.aikg.kgrag_java.service.impl;

import com.aikg.kgrag_java.pojo.dto.neo4j.Neo4jRetrievalDTO;
import com.aikg.kgrag_java.pojo.entity.neo4j.Node;
import com.aikg.kgrag_java.pojo.po.LastOrNextNodeSearchPO;
import com.aikg.kgrag_java.repository.Neo4jRetrievalRepository;
import com.aikg.kgrag_java.service.IEmbeddingService;
import com.aikg.kgrag_java.service.INeo4jService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @ClassName: Neo4jServiceImpl
 * @Description: neo4j检索的服务实现类
 * @Author: 李宗雨
 * @Date: 2025/1/4 11:52
 * @Version: 1.0
 **/
@Slf4j
@Service
public class Neo4jServiceImpl implements INeo4jService {
    @Autowired
    Neo4jRetrievalRepository neo4jRetrievalRepository;

    @Autowired
    IEmbeddingService embeddingService;

    /**
     * neo4j相似度检索
     * @param query: 用户的问题
     * @param projectId: 领域Id
     * @param topK: 设置前k个检索的重要节点
     * @return neo4j检索的结果
     */
    public Neo4jRetrievalDTO neo4jRetrieval(String query, int projectId, int topK){
        // 设置索引名称
        String indexName = "{projectID}_Vector_Index".replace("{projectID}", "" + projectId);
        // 设置待索引的节点类型
        String nodeLabel = "{projectID}_Node".replace("{projectID}", "" + projectId);
        // 设置待索引节点类型的属性
        String embeddingNodeProperty = "Embedding";
        if(!neo4jRetrievalRepository.hasIndex(indexName, nodeLabel, embeddingNodeProperty)){
            // 不存在则创建索引
            neo4jRetrievalRepository.createVectorIndex(indexName,nodeLabel, embeddingNodeProperty);
        }
        Neo4jRetrievalDTO neo4jRetrievalDTO = new Neo4jRetrievalDTO();

        // 记录所有已经检索过的重要节点id，为了防止重复存储节点间的关系
        List<Long> existMainNodeIds = new ArrayList<>();

        // 记录所有已经检索过(重要节点及其邻居节点)的节点id，为了防止重复存储节点
        List<Long> existNodeIds = new ArrayList<>();

        // 找到重要节点并保存重要节点的信息，后续合并用
        List<Node> similarityNodeIdList = this.similaritySearch(query, indexName, 3,0.75f);
//        System.out.println("重要节点:" + similarityNodeIdList);

        // 遍历重要节点
        for(Node node: similarityNodeIdList){

            // 不要忘了增加**本重要节点**
            if(!existNodeIds.contains(node.getId())) {
                neo4jRetrievalDTO.getNodeInfoList().add(node);
                existMainNodeIds.add(node.getId());
                existNodeIds.add(node.getId());
            }
            // 查找上下一跳信息
            Neo4jRetrievalDTO subForMeRetrievalDTO = searchNeighborNode(node, existMainNodeIds, existNodeIds);

            // 合并查到的节点及节点间的关系
            neo4jRetrievalDTO.getNodeInfoList().addAll(subForMeRetrievalDTO.getNodeInfoList());
            neo4jRetrievalDTO.getRelationTripletList().addAll(subForMeRetrievalDTO.getRelationTripletList());
        }
        return neo4jRetrievalDTO;
    }

    /**
     * 按照query进行相似检索
     *
     * @param query: 用户的问题
     * @param indexName: index的名字
     * @param topK: 前k个
     * @param scoreThreshold: 阈值
     * @return
     */
    public List<Node> similaritySearch(String query, String indexName, int topK, float scoreThreshold){
        return neo4jRetrievalRepository.similaritySearch(
                indexName,
                topK,
                embeddingService.embeddingQuery(query),
                scoreThreshold);
    }

    /**
     * 寻找某节点的上一跳和下一跳节点及其属性。
     *  @param node : 重要节点
     *  @param existMainNodeIds : 保存已检索过的重要节点
     *  @param existNodeIds : 保存已检索过的所有节点（包括重要节点及其邻居节点）
     * @return 返回重要节点的所有邻居节点及其关系
     */
    public Neo4jRetrievalDTO searchNeighborNode(Node node, List<Long> existMainNodeIds, List<Long> existNodeIds){
        // 整合所有节点以及节点间的关系三元组
        Neo4jRetrievalDTO forMeRetrievalDTO = new Neo4jRetrievalDTO();

        // 查找下一跳关系及其链接的节点
        List<LastOrNextNodeSearchPO> nodeSearchDTOS = neo4jRetrievalRepository.nextNodeSearch(node.getId());

        // 遍历检索到的信息，将信息整合（节点信息放到节点信息数组nodeInfoList,关系三元组放到relationTripletList中）
        for(LastOrNextNodeSearchPO nodeSearchDTO: nodeSearchDTOS){
            // 判断下一跳节点是否已经检索过，若检索过则跳过，否则加入
            if(!existNodeIds.contains(nodeSearchDTO.getTargetNode().getId())) {
                // 添加尾节点到nodeInfoList中
                forMeRetrievalDTO.getNodeInfoList().add(nodeSearchDTO.getTargetNode());
                existNodeIds.add(nodeSearchDTO.getTargetNode().getId());
            }
                // 添加关系三元组到relationTripletList中
            // 首先判断是否下一个节点是否已经检索过
            if(!existMainNodeIds.contains(nodeSearchDTO.getTargetNode().getId())) {
                Map<String, String> tempTriplet = new HashMap<>();
                tempTriplet.put("头实体", node.getName());
                tempTriplet.put("关系", nodeSearchDTO.getRelationName());
                tempTriplet.put("尾实体", nodeSearchDTO.getTargetNode().getName());
                forMeRetrievalDTO.getRelationTripletList().add(tempTriplet);
            }
        }
        // 查找上一跳关系及其链接的节点
        nodeSearchDTOS = neo4jRetrievalRepository.lastNodeSearch(node.getId());
        for(LastOrNextNodeSearchPO nodeSearchDTO: nodeSearchDTOS){
            // 判断上一跳节点是否已经检索过，若检索过则跳过，否则加入
            if(!existNodeIds.contains(nodeSearchDTO.getHeadNode().getId())) {
                // 添加头节点到nodeInfoList中
                forMeRetrievalDTO.getNodeInfoList().add(nodeSearchDTO.getHeadNode());
                existNodeIds.add(nodeSearchDTO.getHeadNode().getId());
            }

            // 首先判断是否上一个节点是否已经检索过
            if(!existMainNodeIds.contains(nodeSearchDTO.getHeadNode().getId())) {
                // 添加关系三元组到relationTripletList中
                Map<String, String> tempTriplet = new HashMap<>();
                tempTriplet.put("头实体", nodeSearchDTO.getHeadNode().getName());
                tempTriplet.put("关系", nodeSearchDTO.getRelationName());
                tempTriplet.put("尾实体", node.getName());
                forMeRetrievalDTO.getRelationTripletList().add(tempTriplet);
            }
        }
        return forMeRetrievalDTO;
    }

}
